#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：div-align-dg 
@File    ：prep_thyroid_voc_like.py
@IDE     ：PyCharm 
@Author  ：cao xu
@Date    ：2025/8/29 上午11:09 
"""
import os, os.path as osp
from pathlib import Path
from PIL import Image
import xml.etree.ElementTree as ET
import random

# ====== 配置：按需修改 ======
IMAGES_ROOT = "/data/lining/data/Structured_Dataset/Thyroid_Data/Comprehensive_data/picture/images"
YOLO_ROOT   = "/data/lining/data/Structured_Dataset/Thyroid_Data/Comprehensive_data/picture/detect_labels"
VOC_ROOT    = "/data/caoxu/dataset/div-align-dg/voc_labels"
VOC_SETS    = "/data/caoxu/dataset/div-align-dg/voc_sets"
# ====== 医院名单（顶层文件夹名）======
TRAIN_HOSPITALS = {
    "上海十院","上海市一","华西门诊","四川省人民","困难样本","广州市一","徐州市中心",
    "无锡市人民","武汉协和","沈阳医科大","米诺娃","胜利油田","遂宁中心","颐和",
    "成都中科","郑大附一","华西医院","遵义美年","华西-赵婉君","华西-马步云"
}
TEST_HOSPITALS = {
    "上海十院-180例回顾性数据","上海十院-少见癌","华西某院","301桥本结节","公开",
    "陕西肿瘤","无锡某院","绵阳某院","昆明某院","杭州某院"
}
CLASS_KEEP  = {0: "nodule"}                      # 只保留 0=结节，其它类别（如 1=腺体）会被忽略
IMG_EXTS    = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff"}
# ==============================================

def yolo2voc(xc, yc, w, h, W, H):
    x1 = max(0, (xc - w/2) * W); y1 = max(0, (yc - h/2) * H)
    x2 = min(W-1, (xc + w/2) * W); y2 = min(H-1, (yc + h/2) * H)
    return int(round(x1)), int(round(y1)), int(round(x2)), int(round(y2))

def make_xml(img_filename, W, H, objects):
    node = ET.Element("annotation")
    ET.SubElement(node, "folder").text = "VOC_like"
    ET.SubElement(node, "filename").text = img_filename
    size_node = ET.SubElement(node, "size")
    ET.SubElement(size_node, "width").text  = str(W)
    ET.SubElement(size_node, "height").text = str(H)
    ET.SubElement(size_node, "depth").text  = "3"
    ET.SubElement(node, "segmented").text = "0"
    for cls, (x1,y1,x2,y2) in objects:
        obj = ET.SubElement(node, "object")
        ET.SubElement(obj, "name").text = cls
        ET.SubElement(obj, "pose").text = "Unspecified"
        ET.SubElement(obj, "truncated").text = "0"
        ET.SubElement(obj, "difficult").text = "0"
        bb = ET.SubElement(obj, "bndbox")
        ET.SubElement(bb, "xmin").text = str(x1)
        ET.SubElement(bb, "ymin").text = str(y1)
        ET.SubElement(bb, "xmax").text = str(x2)
        ET.SubElement(bb, "ymax").text = str(y2)
    return ET.ElementTree(node)

IMAGES_ROOT = Path(IMAGES_ROOT); YOLO_ROOT = Path(YOLO_ROOT)
VOC_ROOT = Path(VOC_ROOT); VOC_SETS = Path(VOC_SETS)
VOC_ROOT.mkdir(parents=True, exist_ok=True); VOC_SETS.mkdir(parents=True, exist_ok=True)

train_rel_paths, test_rel_paths = [], []
unknown_hospitals = set()

def process_image(img_path: Path):
    rel = img_path.relative_to(IMAGES_ROOT)            # e.g. 301桥本结节/JFJ.../xxx.jpg
    top = rel.parts[0]                                  # 顶层文件夹：医院名
    target_list = None
    if top in TRAIN_HOSPITALS:
        target_list = train_rel_paths
    elif top in TEST_HOSPITALS:
        target_list = test_rel_paths
    else:
        unknown_hospitals.add(top)
        return

    # YOLO 标签同构路径
    yolo_path = (YOLO_ROOT / rel).with_suffix(".txt")
    # 读取图片尺寸
    with Image.open(img_path) as im:
        W, H = im.size

    objects = []
    if yolo_path.exists():
        with open(yolo_path, "r", encoding="utf-8") as f:
            for line in f:
                s = line.strip().split()
                if len(s) != 5:
                    continue
                cid = int(s[0]); xc, yc, w, h = map(float, s[1:])
                if cid not in CLASS_KEEP:  # 丢弃腺体等
                    continue
                box = yolo2voc(xc, yc, w, h, W, H)
                objects.append((CLASS_KEEP[cid], box))

    # 写 VOC XML 到镜像目录
    out_xml = (VOC_ROOT / rel).with_suffix(".xml")
    out_xml.parent.mkdir(parents=True, exist_ok=True)
    xml_tree = make_xml(img_path.name, W, H, objects)
    xml_tree.write(out_xml, encoding="utf-8", xml_declaration=True)

    # 把“相对 images 的相对路径（含扩展名）”写进对应清单
    target_list.append(rel.as_posix())

for p in IMAGES_ROOT.rglob("*"):
    if p.is_file() and p.suffix.lower() in IMG_EXTS:
        process_image(p)

with open(VOC_SETS/"train.txt", "w", encoding="utf-8") as f:
    f.write("\n".join(sorted(train_rel_paths)))
with open(VOC_SETS/"test.txt", "w", encoding="utf-8") as f:
    f.write("\n".join(sorted(test_rel_paths)))

print(f"Train images: {len(train_rel_paths)}")
print(f"Test  images: {len(test_rel_paths)}")
if unknown_hospitals:
    print("Skipped (not in TRAIN/TEST lists):", sorted(unknown_hospitals))
print("VOC labels saved under:", VOC_ROOT)
print("Lists saved under:", VOC_SETS)